package com.kingsoft.dc.khaos.module.spark.preprocess.transform

import com.kingsoft.dc.khaos.KhaosContext
import com.kingsoft.dc.khaos.innertype.Schema
import com.kingsoft.dc.khaos.metadata.{Dependency, KhaosStructField}
import com.kingsoft.dc.khaos.module.spark.metadata.preprocess.transform.{SorterInfo, SorterOption}
import com.kingsoft.dc.khaos.module.spark.util.DataframeUtils
import com.kingsoft.dc.khaos.util.Logging
import org.apache.spark.RangePartitioner
import org.apache.spark.sql.{Column, DataFrame}
import org.json4s.DefaultFormats
import org.json4s.jackson.JsonMethods.parse

import scala.collection.mutable.ArrayBuffer

/**
  * create by yansu on 2019/07/22 19:32
  */
class Sorter extends TransformStrategy with Logging{
  private var sorterInfo: SorterInfo = _
  private val ASC = "asc"
  private val DESC = "desc"

  override def exec(kc: KhaosContext,
                    module_id: String,
                    config: String,
                    dependences: Seq[Dependency],
                    targets: Seq[Dependency]): Seq[(String, DataFrame)] = {
    // 解析config
    implicit val formats = DefaultFormats
    sorterInfo = parse(config, true).extract[SorterInfo]

    var resDataFrame: DataFrame = null
    // 根据依赖取出df
    val sortDataFrame = kc.structData[DataFrame](dependences.head.getEdge())._2

    // 取出排序字段
    val sorterField = sorterInfo.sort_field
    val sortArr = new ArrayBuffer[Column]()


    for (field <- sorterField) {
      val fieldName = field.filed
      val sortType = field.sort_type
      log.info("排序字段 ==> " + fieldName + " 排序类型 ==> " + sortType)

      sortType match {
        case ASC => sortArr += sortDataFrame.col(fieldName).asc
        case DESC => sortArr += sortDataFrame.col(fieldName).desc
      }
    }
    // 根据用户所填字段进行排序
    log.info("开启Sorter算子...")
    resDataFrame = sortDataFrame.repartitionByRange(sortArr:_*).sort(sortArr: _*)
    // 缓存功能
    if ("true".equals(sorterInfo.advanced_options.cache_option.on_off)) {
      resDataFrame = DataframeUtils.cacheDataframe(resDataFrame, this)
    }
    // 分区功能
    if ("true".equals(sorterInfo.advanced_options.partition.on_off)) {
      resDataFrame = DataframeUtils.repartionDataframe(resDataFrame, sorterInfo.advanced_options.partition.nums.toInt, this)
    }
    // top功能
    val top_flag = sorterInfo.advanced_options.top.on_off
    if (top_flag) {
      val topn = sorterInfo.advanced_options.top.topn
      if (topn != null && topn != "" && topn.toInt > 0) {
        try {
          resDataFrame = resDataFrame.limit(topn.toInt)
        } catch {
          case e: IllegalArgumentException =>
            e.printStackTrace()
        }
      } else {
        log.error("未知TOPN: " + topn.toInt)
      }
    }
    addResult(targets.head, resDataFrame)
  }

  //21202F2938212B3E22272626252E434D
  override def schema(kc: KhaosContext,
                      module_id: String,
                      config: String,
                      dependences: Seq[Dependency]) = {

    val schema = if (sorterInfo.extract_fields.isEmpty) {
      kc.schemaChannel.getSchema(dependences.head.getSource())
    } else {
      val schemaList = sorterInfo.extract_fields.map { optInfo =>
        KhaosStructField(optInfo.field, optInfo.data_type)
      }
      new Schema(schemaList)
    }
    schema
  }
}